# -*- coding: utf-8 -*-
"""
-------------------------------------------------
File Name： random_select_branch
Description :
Author : 'li'
date： 2022/6/16
Change Activity:
2022/6/16:
-------------------------------------------------
"""
import copy
import time
from typing import Iterable

import numpy as np

from ml import fetch_element_by_probability
from ml.common.design_pattern.pipline.step_detail import StepDetail


class RandomSelectBranch:
    def __init__(self, transformers_cfg_lst, probabilities_lst, module_registry, save_step_detail=False):
        """

        Args:
            transformers_cfg_lst:
            probabilities_lst:
            module_registry:
            save_step_detail:
        """
        assert len(transformers_cfg_lst) == len(probabilities_lst)
        assert np.array(probabilities_lst).sum() == 1
        self.transformers_cfg_lst = transformers_cfg_lst
        self.probabilities_lst = probabilities_lst
        self.transformers = self._build_transformers()
        self.module_registry = module_registry
        self.save_step_detail = save_step_detail
        self.transformers_branch_lst = self._build_transformers()
        if save_step_detail:
            self.middle_parameters = []  # 保存中间变量，便于debug

    def _build_transformers(self):
        """
        create step objects and combine all branch to transformers.

        """
        transformers_lst = []
        for transformers_cfg in self.transformers_cfg_lst:
            transformers = []
            for step_info in transformers_cfg:
                obj = self._build_step_obj(copy.deepcopy(step_info))
                transformers.append(obj)
            transformers_lst.append(transformers)
        return transformers_lst

    def _handle_one_branch_transformers(self, data, transformers):
        """

        Args:
            transformers:

        Returns:

        """
        start_time, input_data, = None, None
        for trans in transformers:
            if self.save_step_detail:
                start_time = time.time()
                input_data = copy.deepcopy(data)
            data = trans(data)  # handle
            if self.save_step_detail:
                class_name = trans.__class__.__name__
                detail = StepDetail(step_name=class_name, start_time=start_time, end_time=time.time(),
                                    inputs=input_data, outputs=copy.deepcopy(data))
                self.middle_parameters.append(detail)
        return data

    def __call__(self, data):
        """
        执行每一步step

        Args:
            data: 输入数据

        Returns:

        """
        selected_transformers = fetch_element_by_probability(self.transformers_cfg_lst, self.probabilities_lst)
        return self._handle_one_branch_transformers(data, selected_transformers)

    def _build_step_obj(self, step_info):
        """
        构建每一步骤对应的对象

        Args:
            step_info: 来源于pipline的配置文件，根据每一步要执行的信息，新建对应执行对象并返回。

        Returns:
            obj: callable 对象
        """
        step_name = step_info.pop('type')
        _callable = self._get_callable(step_name)
        obj = _callable(**step_info)
        return obj

    def _get_callable(self, callable_name):
        """
        根据 type 的值查找对象模块

        Args:
            callable_name: 模块名

        Returns:

        """
        if isinstance(self.module_registry, Iterable):  # 处理多个registry合并的情况。
            for registry in self.module_registry:
                if registry.contain(callable_name):
                    return registry.get(callable_name)
        else:
            return self.module_registry.get(callable_name)
